# src/data_loader.py

import os
import numpy as np
import torch
from torch.utils.data import Dataset
import pywt

class ECGDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.data_files = [f for f in os.listdir(self.data_dir) if f.endswith('.npy')]

    def __len__(self):
        return len(self.data_files)

    def __getitem__(self, index):
        file_name = self.data_files[index]
        sample_id = os.path.splitext(os.path.basename(file_name))[0]

        file_path = os.path.join(self.data_dir, file_name)
        sample = np.load(file_path, allow_pickle=True).item()
        ecg = sample['ecg']
        label = sample['label']

        # 只返回时域信号
        features = torch.FloatTensor(ecg).unsqueeze(0)  # [1, signal_length]
        label = torch.FloatTensor([label])

        if self.transform:
            features = self.transform(features)

        return features, label, sample_id
